#!/usr/bin/env python3
'''
这种直接收集所有的训练数据，挑选好的训练数据用于训练，训练收敛非常块
比常用的强化学习的方法要块很多，但是存在使用限制的情况。
'''
import gymnasium as gym
from collections import namedtuple
import numpy as np
from tensorboardX import SummaryWriter

import torch
import torch.nn as nn
import torch.optim as optim


HIDDEN_SIZE = 128
BATCH_SIZE = 16
PERCENTILE = 70


class Net(nn.Module):
    def __init__(self, obs_size, hidden_size, n_actions):
        '''

        param obs_size: 观察空间的尺寸
        hidden_size: 训练网络的隐藏层
        n_actions: 动作空间的尺寸
        '''
        super(Net, self).__init__()

        '''
        网络的构成只需要简单的全连接层组成
        '''

        self.net = nn.Sequential(
            nn.Linear(obs_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, n_actions)
        )

    def forward(self, x):
        return self.net(x)


Episode = namedtuple('Episode', field_names=['reward', 'steps'])
EpisodeStep = namedtuple('EpisodeStep', field_names=['observation', 'action'])


def iterate_batches(env, net, batch_size):
    '''
    生成一次训练batch
    '''

    # 本次收集的训练batch、
    # batch的格式是：本次执行一系列动作后知道结束得到的奖励，执行的一系列动作
    batch = []
    #
    episode_reward = 0.0
    # 存储EpisodeStep列表，EpisodeStep类存储每次的环境数据以及该环境对应的执行动作
    episode_steps = []
    # 重置环境
    obs = env.reset()
    # softmax函数，用来将结果转换为概率值
    sm = nn.Softmax(dim=1)
    while True:
        obs_v = torch.FloatTensor([obs]) # 初始环境的观察数据
        act_probs_v = sm(net(obs_v)) # 将观察的数据传入网路。并将结果转换为SoftMax概率
        act_probs = act_probs_v.data.numpy()[0] # 获取本次的推理结果，因为仅传入一个训练数据，所以仅得到一个推理结果
        action = np.random.choice(len(act_probs), p=act_probs) # 从推理的结果中随机获取动作执行
        next_obs, reward, is_done, _ = env.step(action) # 将选择的动作传入环境中继续执行获取执行的结果
        episode_reward += reward # 计算奖励
        episode_steps.append(EpisodeStep(observation=obs, action=action)) # 存储本次的执行动作所处的环境以及动作
        # 是否结束
        if is_done:
            # 结束时将本次的Episode存储在batch中
            # 并重置环境
            batch.append(Episode(reward=episode_reward, steps=episode_steps))
            episode_reward = 0.0
            episode_steps = []
            next_obs = env.reset()
            # 如果已经搜集了足够的数据，则返回收集的数据并等待下一次运行
            if len(batch) == batch_size:
                # 用这种方式返回batch，收集训练数据，可以免除训练数据太多导致内存爆炸的问题
                yield batch
                batch = []
        # 更新环境
        obs = next_obs


def filter_batch(batch, percentile):
    '''
    过滤训练数据 仅留下反馈奖励高的训练数据
    在这篇源码里面，其训练方式是采用了收集奖励高的数据进行训练，通过
    奖励高的训练数据得到好的训练结果

    param batch: 本次训练batch的训练数据，保存本次训练的观察环境以及环境所对应的执行动作

    return: 用于训练的观察环境数据、用于训练的执行动作数据、奖励分数的百分位、奖励分数的中值
    '''

    # 获取奖励列表
    rewards = list(map(lambda s: s.reward, batch))
    # 计算奖励的百分位，通过百分位过滤不合格的训练数据
    reward_bound = np.percentile(rewards, percentile)
    reward_mean = float(np.mean(rewards))

    train_obs = [] # 用于训练的观察环境数据
    train_act = [] # 用于训练的执行动作数据 两者之间一一对应
    for example in batch:
        if example.reward < reward_bound:
            continue
        train_obs.extend(map(lambda step: step.observation, example.steps))
        train_act.extend(map(lambda step: step.action, example.steps))

    train_obs_v = torch.FloatTensor(train_obs)
    train_act_v = torch.LongTensor(train_act)
    return train_obs_v, train_act_v, reward_bound, reward_mean


if __name__ == "__main__":
    # 创建CartPole游戏的环境世界
    env = gym.make("CartPole-v0")
    # env = gym.wrappers.Monitor(env, directory="mon", force=True)
    # 获取观察空间的维度和动作空间的维度
    obs_size = env.observation_space.shape[0]
    n_actions = env.action_space.n

    # 创建训练网络
    net = Net(obs_size, HIDDEN_SIZE, n_actions)
    # 交叉熵损失函数
    objective = nn.CrossEntropyLoss()
    # Adam优化器
    optimizer = optim.Adam(params=net.parameters(), lr=0.01)
    writer = SummaryWriter(comment="-cartpole")

    # 获取本次的训练数据
    for iter_no, batch in enumerate(iterate_batches(env, net, BATCH_SIZE)):
        obs_v, acts_v, reward_b, reward_m = filter_batch(batch, PERCENTILE)
        optimizer.zero_grad()
        action_scores_v = net(obs_v) # 传入执行的环境数据获取动作
        loss_v = objective(action_scores_v, acts_v) # 将动作与推理的结果计算交叉熵 因为推理的出来的动作是一个推理小数值，但是目标的label是已经明确的1， 0值，所以尽管网络一样，但是也可以用于训练，还有一点就是生成训练数据时采用的是随机选择
        loss_v.backward() # 计算梯度
        optimizer.step() # 优化，将梯度更新到参数
        print("%d: loss=%.3f, reward_mean=%.1f, reward_bound=%.1f" % (
            iter_no, loss_v.item(), reward_m, reward_b))
        writer.add_scalar("loss", loss_v.item(), iter_no)
        writer.add_scalar("reward_bound", reward_b, iter_no)
        writer.add_scalar("reward_mean", reward_m, iter_no)
        # 如果获得的奖励累计值能够超过199则表示训练非常成功
        if reward_m > 199:
            print("Solved!")
            break
    writer.close()
